package com.hnit;


import dev.langchain4j.community.model.dashscope.QwenChatModel;
import dev.langchain4j.community.model.dashscope.QwenEmbeddingModel;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.ClassPathDocumentLoader;
import dev.langchain4j.data.document.parser.TextDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentByCharacterSplitter;
import dev.langchain4j.data.document.splitter.DocumentByRegexSplitter;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;

import java.util.List;

public class RAG_App3 {
    public static void main(String[] args) {
        //1.文件加载器有多种实现，这里使用ClassPathDocumentLoader,还有FileDocumentLoader,UrlDocumentLoader等
        //DocumentParser解析器 有多种实现，这里使用TextDocumentParser,还有MarkdownDocumentParser,JsonDocumentParser,XmlDocumentParser等
        Document doc = ClassPathDocumentLoader.loadDocument("rag/a.txt", new TextDocumentParser());
//        System.out.println(doc.text());
        //2.分词器: 决定了文本如何被分割成更小的单元，这些单元可以是单词，句子，段落或其他更细粒度的文本片段。
        // DocumentSplitter splitter = new DocumentByCharacterSplitter(
        //      100,   //每段最长字数
        //      10
        // );
        DocumentByRegexSplitter splitter = new DocumentByRegexSplitter(
                "\\n\\d+\\.\\s+",
                "\\n",
                90,
                1,
                new DocumentByCharacterSplitter(100, 10)
        );
        List<TextSegment> segments = splitter.split(doc);
//        System.out.println(segments);

        //3.向量化器：负责将文本片段转换为数值向量，这些向量可以被机器学习模型用于文本分析，检索和生成。
        // 向量化器有多种实现，这里使用OpenAIEmbeddingService，还有HuggingFaceEmbeddingService等
        QwenEmbeddingModel qwenEmbeddingModel = QwenEmbeddingModel.builder()
                .apiKey(System.getenv("ALI_AI_KEY"))
                .build();

        //向量化
        List<Embedding> embeddings = qwenEmbeddingModel.embedAll(segments).content();

        //4.保存到向量数据库中 此处还是用InMemoryEmbeddingStore ，实际应用中可以使用其他数据库
        //向量数据库有多种实现，这里使用InMemoryEmbeddingStore，还有ElasticsearchEmbeddingStore,RedisEmbeddingStore等
        InMemoryEmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore();
        embeddingStore.addAll(embeddings,segments);


        ////////////////////////////阶段二：检索
        //5.对查询关键词进行向量化
        Response<Embedding> embed = qwenEmbeddingModel.embed("装卸工工作职责");
        //创建请求
        EmbeddingSearchRequest build = EmbeddingSearchRequest.builder()
                .queryEmbedding(embed.content())
                .maxResults(2)
                .build();

        //6.查询
        EmbeddingSearchResult<TextSegment> segmentEmbeddingSearchResult = embeddingStore.search(build);
        segmentEmbeddingSearchResult.matches().forEach( embeddingMatch -> {
            System.out.println("===================================");
            System.out.println(embeddingMatch.score());
            System.out.println(embeddingMatch.embedded().text());
            System.out.println("===================================");
        });

        //////////////////////////阶段三: 检索增强阶段，用户发问->问题向量化->向量数据库查询->返回上下文->prompt->模型生成答案->返回答案
        //大语言模型
        ChatModel chatModel = QwenChatModel.builder()
                .apiKey(System.getenv("ALI_AI_KEY"))
                .modelName("qwen-turbo")
                .build();
        //内容检索器
        ContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStore)
                .embeddingModel(qwenEmbeddingModel)
                .maxResults(5)
                .minScore(0.7)
                .build();
        //为Assistant生成动态代理对象 chat --> 对话内容存储  ChatMemory  --> 聊天
        Assistant assistant = AiServices.builder(Assistant.class)
                .chatModel(chatModel)
                .contentRetriever(contentRetriever)//内容检索
                .build();
        System.out.println("************************");
        System.out.println(assistant.chat("装卸工工作职责"));
    }
    public interface Assistant{
        String chat(String userMessage);
    }
}
